import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import ListedColormap
from sklearn.neighbors import KNeighborsClassifier

data = np.loadtxt('../data/gauss.csv', delimiter=',')
x_train = data[:, 0:2]
y_train = data[:, 2]

# 划分网格
step = 0.02
x_min ,x_max = np.min(x_train[:,0]) - 1, np.max(x_train[:,0]) + 1
y_min ,y_max = np.min(x_train[:,1]) - 1, np.max(x_train[:,1]) + 1
xx, yy = np.meshgrid(np.arange(x_min,x_max,step), np.arange(y_min,y_max,step))
grid_data = np.concatenate([xx.reshape(-1,1), yy.reshape(-1,1)], axis=1)

ks = [1 ,3 ,5]
fig = plt.figure(figsize = (16,4.5))
camp_light = ListedColormap(["royalblue", "lightcoral"])
for i,k in enumerate(ks):
    knn = KNeighborsClassifier(n_neighbors=k)
    knn.fit(x_train, y_train)
    z = knn.predict(grid_data)

    ax = fig.add_subplot(1,3,i+1)
    ax.pcolormesh(xx, yy, z.reshape(xx.shape), cmap=camp_light)
    ax.scatter(x_train[y_train == 0, 0], x_train[y_train == 0, 1], c='blue', marker='o')
    ax.scatter(x_train[y_train == 1, 0], x_train[y_train == 1, 1], c='red', marker='x')

    ax.set(xlabel='x', ylabel='y')
    ax.set_title(f'k={k}')
plt.show()



